import os
import sys
import numpy as np
import matplotlib.pyplot as plt

def match_ios_conv(data):
	out_size = data.shape[0]
	in_size = data.shape[1]
	height = data.shape[2]
	width = data.shape[3]
	for o in range(out_size):
		for h in range(height):
			for w in range(width):
				for i in range(in_size):
					# index = ((o * in_size + i) * height + h) * width + w
					print str(data[o, i, h, w]) + ',',
				if in_size % 4 == 1:
					print '0.0, 0.0, 0.0, ',
				elif in_size % 4 == 2:
					print '0.0, 0.0, ',
				elif in_size % 4 == 3:
					print '0.0, ',
				print
			print
		# print
	print

def match_ios_innerproduct(data, channels, height, width):
	for o in range(data.shape[0]):
		for h in range(height):
			for w in range(width):
				for c in range(channels):
					print str(data[o, (c * height + h) * width + w]) + ', ',
				# print
	print


def print_value(data):
	data = data.reshape(data.size)
	for i in range(data.size):
		print str(data[i]) + ', ',
	print



caffe_root = '/home/chong/caffe/'

sys.path.insert(0, caffe_root + 'python')
import caffe
MODEL_FILE = '/home/chong/caffe/examples/mnist3x3/lenet.prototxt'
PRETRAINED = '/home/chong/caffe/examples/mnist3x3/lenet_iter_20000.caffemodel'

net = caffe.Classifier(MODEL_FILE, PRETRAINED)

IMAGE_FILE = './test1.png'
input_image = caffe.io.load_image(IMAGE_FILE, color = False)

IMAGE_FILE = './test4.bmp'
input_image2 = caffe.io.load_image(IMAGE_FILE, color = False)

print 'input_image:', input_image.shape
# print input_image
caffe.set_mode_gpu()
prediction = net.predict([input_image], oversample = False)

# print 'blobs["data"]: ', net.blobs['data'].data[0]

print 'net.params:', net.params
print 'conv1 shape:', net.params['conv1'][0].data.shape, net.params['conv1'][1].data.shape
print 'conv2 shape:', net.params['conv2'][0].data.shape, net.params['conv2'][1].data.shape
# print 'conv3 shape:', net.params['conv3'][0].data.shape, net.params['conv3'][1].data.shape
# print 'ip1 shape:', net.params['ip1'][0].data.shape, net.params['ip1'][1].data.shape
# print 'ip2 shape:', net.params['ip2'][0].data.shape, net.params['ip2'][1].data.shape

print 'blobs["ip2"]: ', net.blobs['ip2'].data[0]
print 'blobs["prob"]: ', net.blobs['prob'].data[0]
print 'predicted class value: ', prediction[0]
print 'predicted class:', prediction[0].argmax()
# print 'predicted class2:', prediction[1].argmax()
print prediction.shape

# print 'blobs["pool3"]:', dir(net.blobs['pool3'])
# print 'blobs["pool3"]', net.blobs['pool3'].shape
print type(net.blobs)
# print 'blobs["pool3"]', net.blobs['pool3'].data.shape

print 'net params'
for k, v in net.params.items():
	print k, v[0].data.shape, v[1].data.shape
	print str(k) + 'bias: ', v[1].data
	if k == 'ip2':
		print type(v[0].data), v[0].data
		match_ios_innerproduct(v[0].data, 16, 2, 2)
		# print_value(v[0].data[0])
		# match_ios(v[0].data)
		print type(v[1].data), v[1].data

print 'net blobs'
for k, v in net.blobs.items():
	print 'blob', k, v.data.shape
	if k == 'ip2' or k == 'prob':
		print_value(v.data)
		print v.data
	# 	match_ios(v.data)


import time

root_dir = '/home/chong/mnist/t10k/0'

count = 0
# for root, dirs, files in os.walk(root_dir):
# 	for filename in files:
# 		if filename.endswith('png'):
# 			src_path = os.path.join(root, filename)

# 			input_image = caffe.io.load_image(src_path, color = False)
# 			start = time.time()
# 			prediction = net.predict([input_image], oversample = False)[0].argmax()
# 			end = time.time()
# 			print 'use time:', (end - start)*1000, 'ms'

# 			dst_path = os.path.join(root, str(prediction) + '_' + str(count) + '.png')
# 			# os.rename(src_path, dst_path)
# 			print count, src_path, dst_path
# 			count += 1
